#pragma once
#include <memory>
#include <vector>
#include <string>

namespace ast
{
    /* 修改后的文法 */
    // CompUnit -> FuncDefs
    // FuncDefs -> FuncDef FuncDefs?
    // FuncDef -> ("int" | "void") ID "(" Params? ")" Block
    // Params -> Param ("," Params)?
    // Param -> "int" ID
    // Block -> "{" Stmts? "}"
    // Stmts -> Stmt Stmts?
    // Stmt -> Block | ";" | Expr ";" | AssignStmt 
    //         | DefStmt | IfStmt | WhileStmt
    //         | "break" ";" | "continue" ";" | "ret" Expr? ";" 
    // AssignStmt -> ID "=" Expr ";"
    // DefStmt -> "int" ID "=" Expr ";"
    // IfStmt -> "if" "(" Expr ")" Stmt ("else" Stmt)?
    // WhileStmt -> "while" "(" Expr ")" Stmt 
    // Expr -> LOrExpr
    // LOrExpr -> LAndExpr | LOrExpr "||" LAndExpr
    // LAndExpr -> RelExpr | LAndExpr "&&" RelExpr
    // RelExpr -> AddExpr 
    //            | RelExpr ("<" | ">" | "<=" | ">=" | "==" | "!=") AddExpr
    // AddExpr -> MulExpr | AddExpr ("+" | "-") MulExpr
    // MulExpr -> UnaryExpr | MulExpr ("*" | "/" | "%") UnaryExpr
    // UnaryExpr -> PrimaryExpr | ("+" | "-" | "!") UnaryExpr
    // PrimaryExpr -> ID | NUMBER | "(" Expr ")" | CallExpr
    // CallExpr -> ID "(" Args? ")"
    // Args -> Arg ("," Args)?
    // Arg -> Expr

    /* 枚举定义 */
    enum Type { TYPE_INT, TYPE_VOID };

    enum StmtType 
    {
        STMT_BLOCK, STMT_NULL, STMT_EXPR, STMT_ASSIGN, 
        STMT_DEF, STMT_IF, STMT_WHILE,
        STMT_BREAK, STMT_CONT, STMT_RET
    };

    enum ExprType 
    { 
        EXPR_BINARY, EXPR_UNARY, 
        EXPR_ID, EXPR_NUM, EXPR_PAREN, EXPR_CALL 
    };

    enum OpType
    {
        OP_PLUS, OP_MINUS, OP_TIMES, OP_DIV, OP_MOD,
        OP_LT, OP_GT, OP_LE, OP_GE, OP_EQ, OP_NE,
        OP_AND, OP_OR, OP_NOT
    };

    /* AST定义 */ 
    class BaseAST
    {
    public:
        virtual ~BaseAST() = default;
    };

    // FuncDefs -> FuncDef FuncDefs?
    // Params -> Param ("," Params)?
    // Args -> Arg ("," Args)?
    class ListAST : public BaseAST
    {
    public:
        std::vector<std::unique_ptr<BaseAST>> items;
    };

    // CompUnit -> FuncDefs
    class CompUnitAST : public BaseAST
    {
    public:
        std::unique_ptr<BaseAST> func_defs;
    };

    // FuncDef -> ("int" | "void") ID "(" Params? ")" Block
    // Block -> "{" Stmts? "}"
    class FuncDefAST : public BaseAST
    {
    public:
        Type type;
        std::unique_ptr<BaseAST> id;
        std::unique_ptr<BaseAST> params;
        std::unique_ptr<BaseAST> stmts;
    };

    // Param -> "int" ID
    class ParamAST : public BaseAST
    {
    public:
        std::unique_ptr<BaseAST> id;
    };

    // Stmt -> Block | ";" | Expr ";" | AssignStmt 
    //         | DefStmt | IfStmt | WhileStmt
    //         | "break" ";" | "continue" ";" | "ret" Expr? ";" 
    // Block -> "{" Stmts? "}"
    class StmtAST : public BaseAST
    {
    public:
        StmtType stmt_type;
        std::unique_ptr<BaseAST> stmt;
    };

    // AssignStmt -> ID "=" Expr ";"
    class AssignStmtAST : public BaseAST
    {
    public:
        std::unique_ptr<BaseAST> id;
        std::unique_ptr<BaseAST> expr;
    };

    // DefStmt -> "int" ID "=" Expr ";"
    class DefStmtAST : public BaseAST
    {
    public:
        std::unique_ptr<BaseAST> id;
        std::unique_ptr<BaseAST> expr;
    };

    // IfStmt -> "if" "(" Expr ")" Stmt ("else" Stmt)?
    class IfStmtAST : public BaseAST
    {
    public:
        std::unique_ptr<BaseAST> expr;
        std::unique_ptr<BaseAST> then_stmt;
        std::unique_ptr<BaseAST> else_stmt;
    };

    // WhileStmt -> "while" "(" Expr ")" Stmt
    class WhileStmtAST : public BaseAST
    {
    public:
        std::unique_ptr<BaseAST> expr;
        std::unique_ptr<BaseAST> stmt;
    };

    // Expr -> LOrExpr
    // LOrExpr -> LAndExpr 
    // LAndExpr -> RelExpr 
    // RelExpr -> AddExpr 
    // AddExpr -> MulExpr 
    // MulExpr -> UnaryExpr 
    // UnaryExpr -> PrimaryExpr
    // PrimaryExpr -> ID | NUMBER | "(" Expr ")" | AppExpr 
    class ExprAST : public BaseAST
    {
    public:
        ExprType expr_type;
        std::unique_ptr<BaseAST> expr;
    };

    // LOrExpr -> LOrExpr "||" LAndExpr
    // LAndExpr -> LAndExpr "&&" RelExpr
    // RelExpr -> RelExpr ("<" | ">" | "<=" | ">=" | "==" | "!=") AddExpr
    // AddExpr -> AddExpr ("+" | "-") MulExpr
    // MulExpr -> MulExpr ("*" | "/" | "%") UnaryExpr
    class BinaryExprAST : public BaseAST
    {
    public:
        OpType op;
        std::unique_ptr<BaseAST> expr1;
        std::unique_ptr<BaseAST> expr2;
    };

    // UnaryExpr -> PrimaryExpr | ("+" | "-" | "!") UnaryExpr
    class UnaryExprAST : public BaseAST
    {
    public:
        OpType op;
        std::unique_ptr<BaseAST> expr;
    };

    // CallExpr -> ID "(" Args? ")"
    class CallExprAST : public BaseAST
    {
    public:
        std::unique_ptr<BaseAST> id;
        std::unique_ptr<BaseAST> args;
    };

    // Arg -> Expr
    class ArgAST: public BaseAST
    {
    public:
        std::unique_ptr<BaseAST> expr;
    };

    // ID 
    class IdAST : public BaseAST
    {
    public:
        std::string name;
        IdAST(std::string* n) : name(n ? *n : ""){}
        IdAST() = default;
    };

    // NUMBER 
    class NumAST : public BaseAST
    {
    public:
        int val;
        NumAST(int v) : val(v) {}
        NumAST() = default;
    };
}
